from server.agent import model_container
from server.agent.callbacks import Status, CustomAsyncIteratorCallbackHandler
from server.agent.tools.calculate import calculate, CalculatorInput
from server.agent.tools.search_internet import search_internet, SearchInternetInput
from server.agent.custom_template import CustomOutputParser, CustomPromptTemplate
from server.utils import get_ChatOpenAI

import langchain
langchain.verbose = True
from langchain.tools import Tool
from langchain.chains import LLMChain
from langchain.agents import LLMSingleActionAgent, AgentExecutor
from pydantic import BaseModel, Field
import asyncio
import json

from PyCmpltrtok.common import sep

prompt_template = 'Answer the following questions as best you can. If it is in order, you can use some tools appropriately. '\
            'You have access to the following tools:\n\n'\
            '{tools}\n\n'\
            'Use the following format:\n'\
            'Question: the input question you must answer1\n'\
            'Thought: you should always think about what to do and what tools to use.\n'\
            'Action: the action to take, should be one of [{tool_names}]\n'\
            'Action Input: the input to the action\n'\
            'Observation: the result of the action\n'\
            '... (this Thought/Action/Action Input/Observation can be repeated zero or more times)\n'\
            'Thought: I now know the final answer\n'\
            'Final Answer: the final answer to the original input question\n'\
            'Begin!\n\n'\
            'Question: {input}\n\n'\
            'Thought: {agent_scratchpad}\n'

tools = [
    Tool.from_function(
        func=calculate,
        name="calculate",
        description="Useful for when you need to answer questions about simple calculations",
        args_schema=CalculatorInput,
    ),
    Tool.from_function(
        func=search_internet,
        name="search_internet",
        description="Use this tool to use bing search engine to search the internet",
        args_schema=SearchInternetInput,
    ),
]

tool_names = [t.name for t in tools]

if '__main__' == __name__:
    model_container.MODEL = get_ChatOpenAI(
        model_name='qwen-api',
        temperature=0.5,
        max_tokens=4096,
        callbacks=[],
    )
    
    async def agent_chat_iterator(
        query: str,
    ):
        callback = CustomAsyncIteratorCallbackHandler()
        
        model = model_container.MODEL
        
        prompt_template_agent = CustomPromptTemplate(
            template=prompt_template,
            tools=tools,
            input_variables=["input", "intermediate_steps"]
        )
        
        output_parser = CustomOutputParser()
        
        llm_chain = LLMChain(llm=model, prompt=prompt_template_agent)
        
        agent = LLMSingleActionAgent(
                llm_chain=llm_chain,
                output_parser=output_parser,
                stop=["\nObservation:", "Observation"],
                allowed_tools=tool_names,
            )
        agent_executor = AgentExecutor.from_agent_and_tools(
            agent=agent,
            tools=tools,
            verbose=True,
        )
        
        async def task_fn():
            await agent_executor.acall(query, callbacks=[callback], include_run_info=True)
            callback.done.set()
            
        task = asyncio.create_task(task_fn())
        
        async for chunk in callback.aiter():
            tools_use = []
            # Use server-sent-events to stream the response
            data = json.loads(chunk)
            if data["status"] == Status.start or data["status"] == Status.complete:
                continue
            elif data["status"] == Status.error:
                tools_use.append("\n```\n")
                tools_use.append("工具名称: " + data["tool_name"])
                tools_use.append("工具状态: " + "调用失败")
                tools_use.append("错误信息: " + data["error"])
                tools_use.append("重新开始尝试")
                tools_use.append("\n```\n")
                yield json.dumps({"tools": tools_use}, ensure_ascii=False)
            elif data["status"] == Status.tool_finish:
                tools_use.append("\n```\n")
                tools_use.append("工具名称: " + data["tool_name"])
                tools_use.append("工具状态: " + "调用成功")
                tools_use.append("工具输入: " + data["input_str"])
                tools_use.append("工具输出: " + data["output_str"])
                tools_use.append("\n```\n")
                yield json.dumps({"tools": tools_use}, ensure_ascii=False)
            elif data["status"] == Status.agent_finish:
                yield json.dumps({"final_answer": data["final_answer"]}, ensure_ascii=False)
            else:
                yield json.dumps({"answer": data["llm_token"]}, ensure_ascii=False)
        
        await task
        
    async def main():
        while True:
            sep()
            print('Your query?: ')
            query = input().strip()
            async for xjson in agent_chat_iterator(query):
                xdict = json.loads(xjson)
                print(xdict)
                
    asyncio.run(main())
    